package com.example.mcp.client;

import com.example.mcp.server.model.McpInitializeRequest;
import com.example.mcp.server.model.McpInitializeResponse;
import com.example.mcp.server.model.McpListToolsResponse;
import com.example.mcp.server.model.McpToolCallRequest;
import com.example.mcp.server.model.McpToolCallResponse;
import com.fasterxml.jackson.databind.ObjectMapper;
import java.util.HashMap;
import java.util.Map;
import java.util.UUID;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpMethod;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.client.RestTemplate;

public class McpClient {

    private final String baseUrl;
    private final RestTemplate restTemplate;
    private final ObjectMapper objectMapper;
    private String sessionId;

    public McpClient(String baseUrl) {
        this.baseUrl = baseUrl;
        this.restTemplate = new RestTemplate();
        this.objectMapper = new ObjectMapper();
    }

    public void connect() {
        String sessionUrl = baseUrl + "/mcp/session";
        ResponseEntity<Map> response = restTemplate.postForEntity(sessionUrl, null, Map.class);
        this.sessionId = (String) response.getBody().get("sessionId");
        System.out.println("Connected to MCP server with session ID: " + sessionId);
    }

    public void initialize() {
        String initUrl = baseUrl + "/mcp/initialize";

        McpInitializeRequest request = new McpInitializeRequest();
        request.setId(UUID.randomUUID().toString());
        request.setProtocolVersion("2024-11-05");

        McpInitializeRequest.ClientCapabilities capabilities = new McpInitializeRequest.ClientCapabilities();
        capabilities.setExperimental(null);
        capabilities.setSampling(null);
        request.setCapabilities(capabilities);

        McpInitializeRequest.ClientInfo clientInfo = new McpInitializeRequest.ClientInfo();
        clientInfo.setName("Spring AI MCP Client");
        clientInfo.setVersion("1.0.0");
        request.setClientInfo(clientInfo);

        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_JSON);

        HttpEntity<McpInitializeRequest> entity = new HttpEntity<>(request, headers);
        ResponseEntity<McpInitializeResponse> response = restTemplate.postForEntity(initUrl, entity, McpInitializeResponse.class);

        System.out.println("Initialized MCP server: " + response.getBody().getServerInfo().getName());
    }

    public void listTools() {
        String toolsUrl = baseUrl + "/mcp/tools";

        HttpHeaders headers = new HttpHeaders();
        headers.set("X-Session-ID", sessionId);

        HttpEntity<String> entity = new HttpEntity<>(headers);
        ResponseEntity<McpListToolsResponse> response = restTemplate.exchange(toolsUrl, HttpMethod.GET, entity, McpListToolsResponse.class);

        System.out.println("Available tools:");
        response.getBody().getTools().forEach(tool -> {
            System.out.println("- " + tool.getName() + ": " + tool.getDescription());
        });
    }

    public Object callTool(String toolName, Map<String, Object> arguments) {
        String callUrl = baseUrl + "/mcp/tools/call";

        McpToolCallRequest request = new McpToolCallRequest();
        request.setId(UUID.randomUUID().toString());
        request.setName(toolName);
        request.setArguments(arguments);

        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_JSON);
        headers.set("X-Session-ID", sessionId);

        HttpEntity<McpToolCallRequest> entity = new HttpEntity<>(request, headers);
        ResponseEntity<McpToolCallResponse> response = restTemplate.postForEntity(callUrl, entity, McpToolCallResponse.class);

        return response.getBody().getContent();
    }

    public void disconnect() {
        if (sessionId != null) {
            String disconnectUrl = baseUrl + "/mcp/session/" + sessionId;
            restTemplate.delete(disconnectUrl);
            System.out.println("Disconnected from MCP server");
        }
    }

    public static void main(String[] args) {
        McpClient client = new McpClient("http://localhost:8080");

        try {
            client.connect();
            client.initialize();
            client.listTools();

            // Test weather tool
            Map<String, Object> weatherArgs = new HashMap<>();
            weatherArgs.put("location", "Beijing, China");
            weatherArgs.put("unit", "celsius");

            Object weatherResult = client.callTool("get_weather", weatherArgs);
            System.out.println("Weather result: " + weatherResult);

            // Test calculator tool
//            Map<String, Object> calcArgs = new HashMap<>();
//            calcArgs.put("expression", "15 * 3 + 7");
//
//            Object calcResult = client.callTool("calculate", calcArgs);
//            System.out.println("Calculator result: " + calcResult);
//
            // Test AI assistant
            Map<String, Object> aiArgs = new HashMap<>();
            aiArgs.put("query", "What is Spring AI and how does it work?");

            Object aiResult = client.callTool("ai_assistant", aiArgs);
            System.out.println("AI response: " + aiResult);

        } finally {
            client.disconnect();
        }
    }
}